import os
import re
import chardet

def detect_encoding(file_path):
    """检测文件编码"""
    with open(file_path, 'rb') as file:
        raw_data = file.read()
        result = chardet.detect(raw_data)
        encoding = result['encoding']
        # 如果检测到的编码是 EUC-TW，替换为 utf-8
        if encoding == 'EUC-TW':
            encoding = 'utf-8'
        return encoding


def clean_data(input_file):
    """清洗 .srt 文件，去除多余信息和无意义字幕段，返回清洗后的字幕内容。"""
    # 尝试检测文件编码
    encoding = detect_encoding(input_file)
    if not encoding:
        encoding = 'utf-8'  # 默认使用 utf-8

    try:
        with open(input_file, "r", encoding=encoding) as file:
            content = file.read()
    except UnicodeDecodeError:
        print(f"尝试使用 {encoding} 编码读取文件{input_file}失败，尝试使用 'latin1' 编码")
        with open(input_file, "r", encoding='latin1') as file:
            content = file.read()

    # 分割每个字幕段落（字幕编号 + 时间范围 + 文本内容）
    subtitle_blocks = re.split(r"\n\s*\n", content.strip())
    cleaned_blocks = []

    for block in subtitle_blocks:
        lines = block.split("\n")

        # 检查是否是标准字幕段
        if len(lines) >= 3 and re.match(r"\d+", lines[0]) and re.match(
                r"\d{2}:\d{2}:\d{2},\d{3} --> \d{2}:\d{2}:\d{2},\d{3}", lines[1]):
            # 清除包含 "Improved by"、"Resynced by"、"Enjoy the movie" 等内容的行
            text_lines = lines[2:]
            cleaned_text = "\n".join(
                [line for line in text_lines if
                 not re.search(r"(Improved by|Resynced by|Enjoy the movie)", line, re.IGNORECASE)]
            )
            if cleaned_text.strip():  # 确保清理后字幕内容不为空
                lines[2:] = cleaned_text.split("\n")
                cleaned_blocks.append("\n".join(lines))

    return "\n\n".join(cleaned_blocks)


def process_files_in_directory(directory, output_file):
    # 在开始时打开输出文件，并清空已有内容
    with open(output_file, 'w', encoding='utf-8') as outfile:
        for filename in os.listdir(directory):
            file_path = os.path.join(directory, filename)
            # 检查文件是否是 .srt 文件
            if os.path.isfile(file_path) and file_path.endswith('.srt'):
                cleaned_content = clean_data(file_path)
                # 将清洗后的字幕内容写入文件，避免文件内容重复写入
                if cleaned_content.strip():
                    outfile.write(cleaned_content + '\n')
y = 2022
for i in range(1, 13):
    directory_to_clean = fr'D:\subtitle\moves\{y}\{i}月'

    if i < 10:
        output_file_path = fr'D:\subtitle\moves\txt\subtitle{y}0{i}.txt'
    else:
        output_file_path = fr'D:\subtitle\moves\txt\subtitle{y}{i}.txt'
    process_files_in_directory(directory_to_clean, output_file_path)
    print(f"字幕文件{i}清洗完成。")

# # 使用示例
# directory_to_clean = r'D:\subtitle\moves\2024\1月'
# output_file_path = r'D:\subtitle\moves\txt\subtitle202401.txt'
# # directory_to_clean = r'D:\subtitle\bring_goods\24'
# # output_file_path = r'D:\subtitle\bring_goods\txt\subtitle24.txt'
#
# # 调用函数进行处理
# process_files_in_directory(directory_to_clean, output_file_path)
#
# print("字幕文件清洗完成。")
